#include <iostream>
#include <cstdlib>
#include <cstdio>
#include <vector>
#include <set>
#include "rigelprint.h"

int main(int argc, char *argv[])
{
  if(argc != 9) {
    fprintf(stderr, "Usage: %s <nx> <ny> <nz> <compute mean> <compute std dev> <comm mean> <comm std dev> <outfilename>\n", argv[0]);
    exit(1);
  }

  int nx = atoi(argv[1]), ny = atoi(argv[2]), nz = atoi(argv[3]);
  float compmean = (float)atof(argv[4]), compsd = (float)atof(argv[5]);
  float commmean = (float)atof(argv[6]), commsd = (float)atof(argv[7]);

  FILE *out = fopen(argv[8], "w");
  if(out == NULL) {
    fprintf(stderr, "Error: could not open %s\n", argv[8]);
    exit(1);
  }

  fprintf(out, "digraph fft_%d_%d_%d {\n", nx, ny, nz);
  fprintf(out, "node [v_delay_mean=0, v_delay_sd=0]\n");
  for(int i = 0; i < nx; i++) {
    for(int j = 0; j < ny; j++) {
      for(int k = 0; k < nz; k++) {
        fprintf(out, "in_%d_%d_%d\n", i, j, k);
        fprintf(out, "out_%d_%d_%d\n", i, j, k);
      }
    }
  }
  fprintf(out, "done\n");

  std::vector<int> finderxy, finderyz;
  finderxy.resize(nx*ny*nz);
  finderyz.resize(nx*ny*nz);
  
  std::vector<std::vector<std::set<int> > > l; //Outer dimension is FFT dimension, middle dimension is the lists of all compute nodes within that stage, inner dimension (ints) is the list of data required.
  l.resize(3);
  l[0].resize(ny*nz);
  l[1].resize(nx*nz);
  l[2].resize(nx*ny);
  for(int i = 0; i < nx*ny*nz; i++) {
    int xbin = i / nx;
    int ybin = ((i/(ny*nz))*nx) + (i % nx);
    int zbin = i % (nx*ny);
    l[0][xbin].insert(i);
    l[1][ybin].insert(i);
    finderxy[i] = ybin;
    l[2][zbin].insert(i);
    finderyz[i] = zbin;
  }
  fprintf(out, "node [v_delay_mean=%f, v_delay_sd=%f]\n", compmean, compsd);
  for(int i = 0; i < ((ny*nz)+(nx*nz)+(nx*ny)); i++) {
    fprintf(out, "fft%d\n", i);
  }

  fprintf(out, "edge [e_delay_mean=%f, e_delay_sd=%f]\n", commmean, commsd);

  //Now write out edges for all the connections between nodes with matching data.
  for(int j = 0; j < l[0].size(); j++) {
    for(std::set<int>::const_iterator it = l[0][j].begin(), end = l[0][j].end(); it != end; ++it) {
      int k = finderxy[*it];
      fprintf(out, "fft%d", j);
      fprintf(out, " -> fft%d\n", l[0].size()+k);
    }
  }
  
  for(int j = 0; j < l[1].size(); j++) {
    for(std::set<int>::const_iterator it = l[1][j].begin(), end = l[1][j].end(); it != end; ++it) {
      int k = finderyz[*it];
      fprintf(out, "fft%d", l[0].size()+j);
      fprintf(out, " -> fft%d\n", l[0].size()+l[1].size()+k);
    }
  }

  //Now write out edges for all the connections for input and output nodes
  fprintf(out, "edge [e_delay_mean=%f, e_delay_sd=%f]\n", commmean, commsd);
  
  for(int i = 0; i < l[0].size(); i++) {
    for(std::set<int>::const_iterator it = l[0][i].begin(), end = l[0][i].end(); it != end; ++it) {
      int a = *it;
      int x = a % nx;
      int y = (a / nx) % ny;
      int z = a / (nx*ny);
      fprintf(out, "in_%d_%d_%d -> fft%d\n", x, y, z, i);
    }
  }

  for(int i = 0; i < l[2].size(); i++) {
    for(std::set<int>::const_iterator it = l[2][i].begin(), end = l[2][i].end(); it != end; ++it) {
      int a = *it;
      int x = a % nx;
      int y = (a / nx) % ny;
      int z = a / (nx*ny);
      fprintf(out, "fft%d", l[0].size()+l[1].size()+i);
      fprintf(out, " -> out_%d_%d_%d\n", x, y, z);
    }
  }

  //Want edges from all outputs to "done" vertex.
  fprintf(out, "edge [e_delay_mean=0, e_delay_sd=0]\n");
  
  for(int i = 0; i < nx; i++) {
    for(int j = 0; j < ny; j++) {
      for(int k = 0; k < nz; k++) {
        fprintf(out, "out_%d_%d_%d -> done\n", i, j, k);
      }
    }
  }
  
  fprintf(out, "}\n");
  fclose(out);
  return 0;
}
